#include "ragwindow.h"
#include <QFileDialog>
#include <QMessageBox>

RAGWindow::RAGWindow(QWidget *parent)
    : QMainWindow(parent)
{
    setupUI();

    // 初始化嵌入管理器
    embeddingManager = new EmbeddingManager(this);
    if (!embeddingManager->initialize()) {
        QMessageBox::critical(this, "错误", "初始化嵌入模型失败！");
        close();
    }

    showStatus("就绪");
}

void RAGWindow::setupUI()
{
    // 主窗口设置
    setWindowTitle("RAG-资料库信息提取");


    // 中央部件
    centralWidget = new QWidget(this);
    setCentralWidget(centralWidget);
    mainLayout = new QVBoxLayout(centralWidget);
    mainLayout->setContentsMargins(30, 30, 30, 30);
    mainLayout->setSpacing(20);

    // 参数设置区域
    QGroupBox *paramGroup = new QGroupBox("参数配置", centralWidget);
    paramGroup->setStyleSheet("QGroupBox { font-weight: bold; }");
    QFormLayout *paramLayout = new QFormLayout(paramGroup);
    paramLayout->setSpacing(10);

    maxChunkSpin = new QSpinBox();
    maxChunkSpin->setRange(50, 1000);
    maxChunkSpin->setValue(200);
    paramLayout->addRow("最大块长度:", maxChunkSpin);

    minChunkSpin = new QSpinBox();
    minChunkSpin->setRange(10, 500);
    minChunkSpin->setValue(20);
    paramLayout->addRow("最小块长度:", minChunkSpin);

    topKSpin = new QSpinBox();
    topKSpin->setRange(1, 20);
    topKSpin->setValue(3);
    paramLayout->addRow("返回结果数:", topKSpin);

    // 文件处理按钮
    processFileBtn = new QPushButton("选择文件");
    processFileBtn->setStyleSheet(
        "QPushButton {"
        "   background-color: #6495ED;"
        "   color: white;"
        "   border: none;"
        "   padding: 12px 24px;"
        "   border-radius: 6px;"
        "   font-size: 16px;"
        "   min-width: 120px;"
        "}"
        "QPushButton:hover {"
        "   background-color: #4169E1;"
        "}"
        "QPushButton:pressed {"
        "   background-color: #1E90FF;"
        "}");
    processFileBtn->setSizePolicy(QSizePolicy::Fixed, QSizePolicy::Fixed);

    // 查询区域
    QGroupBox *queryGroup=new QGroupBox("查询",centralWidget);
    queryGroup->setStyleSheet("QGroupBox { font-weight: bold; }");
    QFormLayout *queryLayout=new QFormLayout(queryGroup);

    queryInput = new QTextEdit();
    queryInput->setPlaceholderText("输入您的问题...");
    queryInput->setMaximumHeight(100);
    queryLayout->addWidget(queryInput);

    queryBtn = new QPushButton("查询");
    queryBtn->setStyleSheet(
        "QPushButton {"
        "   background-color: #6495ED;"
        "   color: white;"
        "   border: none;"
        "   padding: 12px 24px;"
        "   border-radius: 6px;"
        "   font-size: 16px;"
        "   min-width: 120px;"
        "}"
        "QPushButton:hover {"
        "   background-color: #4169E1;"
        "}"
        "QPushButton:pressed {"
        "   background-color: #1E90FF;"
        "}");
    queryBtn->setSizePolicy(QSizePolicy::Fixed, QSizePolicy::Fixed);


    // 结果列表
    QGroupBox *resultsGroup=new QGroupBox("结果",centralWidget);
    resultsGroup->setStyleSheet("QGroupBox { font-weight: bold; }");
    QFormLayout *resultLayout=new QFormLayout(resultsGroup);

    resultsList = new QListWidget();
    resultLayout->addWidget(resultsList);

    statusLabel = new QLabel();
    statusLabel->setAlignment(Qt::AlignCenter);
    resultLayout->addWidget(statusLabel);


    // 控制按钮
    clearBtn = new QPushButton("清空数据库");
    clearBtn->setStyleSheet(
        "QPushButton {"
        "   background-color: #6495ED;"
        "   color: white;"
        "   border: none;"
        "   padding: 12px 24px;"
        "   border-radius: 6px;"
        "   font-size: 16px;"
        "   min-width: 120px;"
        "}"
        "QPushButton:hover {"
        "   background-color: #4169E1;"
        "}"
        "QPushButton:pressed {"
        "   background-color: #1E90FF;"
        "}");
    clearBtn->setSizePolicy(QSizePolicy::Fixed, QSizePolicy::Fixed);

    // 状态栏


    // 布局
    mainLayout->addWidget(paramGroup);

    mainLayout->addWidget(processFileBtn, 0, Qt::AlignHCenter);


    mainLayout->addWidget(queryGroup);
    mainLayout->addWidget(queryBtn, 0, Qt::AlignHCenter);
    mainLayout->addWidget(resultsGroup);
    mainLayout->addStretch();
    mainLayout->addWidget(clearBtn, 0, Qt::AlignHCenter);
    //mainLayout->addWidget(statusLabel);

    // 连接信号槽
    connect(processFileBtn, &QPushButton::clicked, this, &RAGWindow::processFile);
    connect(queryBtn, &QPushButton::clicked, this, &RAGWindow::queryDatabase);
    connect(clearBtn, &QPushButton::clicked, this, &RAGWindow::clearDatabase);

    resize(1360, 765);
}

void RAGWindow::processFile()
{
    QString filePath = QFileDialog::getOpenFileName(this, "选择文件",
                                                  QDir::homePath(),
                                                  "文本文件 (*.txt *.md *.pdf *.docx)");

    if (filePath.isEmpty()) return;

    showStatus("正在处理文件...");

    int maxChunk = maxChunkSpin->value();
    int minChunk = minChunkSpin->value();

    QStringList files = {filePath};
    embeddingManager->processDocuments(files, maxChunk, minChunk);

    showStatus(QString("文件处理完成: %1").arg(QFileInfo(filePath).fileName()));
}

void RAGWindow::queryDatabase()
{
    QString question = queryInput->toPlainText().trimmed();
    if (question.isEmpty()) {
        QMessageBox::warning(this, "提示", "请输入查询问题");
        return;
    }

    showStatus("正在查询...");

    int topK = topKSpin->value();
    QStringList results = embeddingManager->queryRAG(question, topK);

    resultsList->clear();
    if (results.isEmpty()) {
        resultsList->addItem("未找到相关结果");
    } else {
        int num = 1;
        for(const QString &res : results) {
            resultsList->addItem(QString("结果%1:\n%2").arg(num).arg(res));
            num++;
        }
    }

    showStatus(QString("找到 %1 个结果").arg(results.size()));
}

void RAGWindow::clearDatabase()
{
    QMessageBox::StandardButton reply;
    reply = QMessageBox::question(this, "确认", "确定要清空向量数据库吗？",
                                  QMessageBox::Yes|QMessageBox::No);

    if (reply == QMessageBox::Yes) {
        embeddingManager->resetVectorDatabase();
        showStatus("数据库已清空");
    }
}

void RAGWindow::showStatus(const QString &message)
{
    statusLabel->setText(message);
    QCoreApplication::processEvents();
}
